import cv2
import numpy as np
import glob
import yaml
import os

def init_cam(cam_num=0):
    #if os.path.isfile('calibration_matrix.yaml'):
    #    print('have been done, exit')
    #    return 0
    w=10
    h=7
    # 找棋盘格角点
    # 设置寻找亚像素角点的参数，采用的停止准则是最大循环次数30和最大误差容限0.001
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)  # 阈值
    # 世界坐标系中的棋盘格点,例如(0,0,0), (1,0,0), (2,0,0) ....,(8,5,0)，去掉Z坐标，记为二维矩阵
    objp = np.zeros((w * h, 3), np.float32)
    objp[:, :2] = np.mgrid[0:w, 0:h].T.reshape(-1, 2)
    objp = objp * 20
    
    # 储存棋盘格角点的世界坐标和图像坐标对
    objpoints = []  # 在世界坐标系中的三维点
    imgpoints = []  # 在图像平面的二维点
    # 加载pic文件夹下所有的jpg图像
    #images = glob.glob(image_path)  # 拍摄的十几张棋盘图片所在目录
    
    for i in range(20):
        cap = cv2.VideoCapture(cam_num)
        ret, img = cap.read()
        if img is None:
            print(f"无法使用图片: {i}")
            continue
        
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        # 找到棋盘格角点
        ret, corners = cv2.findChessboardCorners(gray, (w, h), None)
        # 如果找到足够点对，将其存储起来
        if ret:
            print(f"使用角点的图片索引: {i}")
            # 在原角点的基础上寻找亚像素角点
            cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
            # 追加进入世界三维点和平面二维点中
            objpoints.append(objp)
            imgpoints.append(corners)
    
    
    # 标定
    print('正在计算')
    # 标定
    ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(objpoints, imgpoints, gray.shape[::-1], None, None)
    
    print(f"标定结果: {ret}")
    print(f"内参数矩阵:\n{mtx}")
    print(f"畸变系数:\n{dist}")
    print(f"旋转(向量)外参:\n{rvecs}")
    print(f"平移(向量)外参:\n{tvecs}")
    
    # 获取优化后的相机矩阵
    u, v = img.shape[:2]
    newcameramtx, roi = cv2.getOptimalNewCameraMatrix(mtx, dist, (u, v), 0, (u, v))
    print(f"优化后的相机矩阵: {newcameramtx}")
    
    # 保存内参矩阵和畸变系数到yaml文件
    file_path = "./calibration_matrix.yaml"
    mtx_yaml = mtx.tolist()
    dist_yaml = dist.tolist()
    data = {"information": "Camera calibration parameters", "camera_matrix": mtx_yaml, "dist_coeff": dist_yaml}
    with open(file_path, "w") as file:
        yaml.dump(data, file)
    print(f"内参和畸变系数已保存到文件: {file_path}")
    
    # 计算重投影误差
    total_error = 0
    for i in range(len(objpoints)):
        imgpoints2, _ = cv2.projectPoints(objpoints[i], rvecs[i], tvecs[i], mtx, dist)
        error = cv2.norm(imgpoints[i], imgpoints2, cv2.NORM_L2) / len(imgpoints2)
        total_error += error
        print(f"第 {i} 张图片的重投影误差: {error}")
    
    mean_error = total_error / len(objpoints)
    print(f"平均重投影误差: {mean_error}")
    
    return 0

def setup_cam(chessboard_size = (10, 7), square_size = 20, cam_num=0):
    '''
    chessboard_size棋盘格内角点数量
    square_size棋盘格每个方格的实际大小
    '''

    def load_calibration_data(file_path):
        if not os.path.exists(file_path):
            print(f"内参文件 {file_path} 不存在。")
            return None, None
        try:
            with open(file_path, 'r') as file:
                data = yaml.safe_load(file)
            mtx = np.array(data['camera_matrix'])
            dist = np.array(data['dist_coeff'])
            return mtx, dist
        except Exception as e:
            print(f"加载内参文件 {file_path} 时出错: {e}")
            return None, None


    # 保存外参矩阵
    def save_extrinsic_data(file_path, rvec, tvec):
        data = {"rotation_vector": rvec.tolist(), "translation_vector": tvec.tolist()}
        try:
            with open(file_path, "w") as file:
                yaml.dump(data, file)
            print(f"外参矩阵已保存到文件: {file_path}")
        except Exception as e:
            print(f"保存外参文件 {file_path} 时出错: {e}")


    # 设置文件路径
    calibration_file = "calibration_matrix.yaml"
    target_folder = "target"
    output_file = "extrinsic_matrix.yaml"

    # 加载内参矩阵和畸变系数
    mtx, dist = load_calibration_data(calibration_file)
    if mtx is None or dist is None:
        print("无法加载内参矩阵和畸变系数，程序终止。")
        exit(1)

    # 世界坐标系中的棋盘格点
    objp = np.zeros((chessboard_size[0] * chessboard_size[1], 3), np.float32)
    objp[:, :2] = np.mgrid[0:chessboard_size[0], 0:chessboard_size[1]].T.reshape(-1, 2)
    objp = objp * square_size


    # 使用通配符获取 target 文件夹中的图片
    #image_paths = glob.glob(os.path.join(target_folder, '*'))
    #if not image_paths:
    #    print(f"未在 {target_folder} 文件夹中找到图片。")
    #    exit(1)


    for i in range(5):
        # 读取图片
        cap = cv2.VideoCapture(cam_num)
        ret, img = cap.read()
        if img is None:
            continue

        # 不缩放图片
        # width = 800
        # r = width / float(img.shape[1])
        # dim = (width, int(img.shape[0] * r))
        # img = cv2.resize(img, dim, interpolation=cv2.INTER_AREA)

        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        ret, corners = cv2.findChessboardCorners(gray, chessboard_size, None)

        if ret:
            # 提高角点检测精度
            criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
            cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)

            # 使用 solvePnP 计算外参
            ret, rvec, tvec = cv2.solvePnP(objp, corners, mtx, dist)

            if ret:
                # 将旋转向量转换为旋转矩阵
                rmat, _ = cv2.Rodrigues(rvec)

                print(f"图片: {i}")
                print("旋转矩阵:")
                print(rmat)
                print("平移向量:")
                print(tvec)

                # 保存外参矩阵
                save_extrinsic_data(output_file, rvec, tvec)

                # 计算重投影误差
                reprojected_corners, _ = cv2.projectPoints(objp, rvec, tvec, mtx, dist)
                reprojected_corners = np.squeeze(reprojected_corners)
                corners = np.squeeze(corners)  # 调整 corners 的形状
                error = cv2.norm(corners, reprojected_corners, cv2.NORM_L2) / len(corners)
                print(f"重投影误差: {error}")

                # 可视化
                #img_with_corners = cv2.drawChessboardCorners(img, chessboard_size, corners.reshape(-1, 1, 2), ret)
                #cv2.imshow('Chessboard Corners', img_with_corners)
                #cv2.waitKey(0)
            else:
                print(f"solvePnP 失败: {i}")
                print(f"objp 形状: {objp.shape}, 值: {objp}")
                print(f"corners 形状: {corners.shape}, 值: {corners}")
        else:
            print(f"未找到棋盘格角点: {i}")


    return 0


if __name__ == '__main__':
    init_cam()
    setup_cam()